import re
from datetime import datetime
from typing import List, Sequence

from llama_index.core import PromptTemplate, Settings, QueryBundle
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from llama_index.core.base.response.schema import RESPONSE_TYPE
from llama_index.core.callbacks import EventPayload, CBEventType
from llama_index.core.prompts import PromptType
from llama_index.core.query_engine import KnowledgeGraphQueryEngine
from llama_index.core.response_synthesizers import (
    BaseSynthesizer,
    get_response_synthesizer,
)
from llama_index.core.schema import NodeWithScore, logger, TextNode
from llama_index.core.utils import print_text

from server.config import vn

# Prompt
DEFAULT_NEBULAGRAPH_NL2CYPHER_PROMPT_TMPL = """
    Generate NebulaGraph query from natural language.
    Use only the provided relationship types and properties in the schema.
    Do not use any other relationship types or properties that are not provided.
    Schema:
    ---
    {schema}
    ---
    Note: NebulaGraph speaks a dialect of Cypher, comparing to standard Cypher:

    1. it uses double equals sign for comparison: `==` rather than `=`
    2. it needs explicit label specification when referring to node properties, i.e.
    v is a variable of a node, and we know its label is Foo, v.`foo`.name is correct
    while v.name is not.

    For example, see this diff between standard and NebulaGraph Cypher dialect:
    ```diff
    < MATCH (p:person)-[:directed]->(m:movie) WHERE m.name = 'The Godfather'
    < RETURN p.name;
    ---
    > MATCH (p:`person`)-[:directed]->(m:`movie`) WHERE m.`movie`.`name` == 'The Godfather'
    > RETURN p.`person`.`name`;
    ```

    Question: {query_str}

    NebulaGraph Cypher dialect query:
    """

DEFAULT_NEBULAGRAPH_NL2CYPHER_PROMPT_TMPL1 = """
## Role
你是一个图数据库专家，请根据用户的问题生成NebulaGraph图数据库的Cypher查询语句。

## Goals
使用Schema中提供的图数据库的关系类型、属性和Context中的上下文，根据用户的问题，请生成NebulaGraph图数据库的cypher查询语句，生成的Cypher语句请严格遵守Attention中的规则。

## Schema
---
{schema}
--- 

## Context
---
当前时间是：{current_time}，如果要查询某一时间区间内的数据，请根据当前时间计算查询的开始时间和结束时间，使用计算出的开始结束时间来作为查询条件。
NebulaGraph中datetime类型支持的属性名称包括year、month、day、hour、minute和second。例如可以使用db.database.create_time.year获取年份，返回的年份是数值类型。
NebulaGraph中可以使用toString()函数将数值类型转为字符串类型。
{context}
--- 
   
## Attention
1. 仅使用Schema中提供的关系类型和属性，不要使用任何其他未提供的关系类型或属性。
2. NebulaGraph使用Cypher方言，与标准Cypher相比有一些区别。例如，请查看标准Cypher和NebulaGraph Cypher方言之间的差异:
    ```diff
    # 标准Cypher
    < MATCH (p:person)-[d:directed]->(m:movie) WHERE m.name = 'The Godfather' and d.time = 2024
    < RETURN p.name;
    ---
    # NebulaGraph Cypher方言
    > MATCH (p:`person`)-[d:directed]->(m:`movie`) WHERE m.`movie`.`name` == 'The Godfather' and d.time = 2024
    > RETURN p.`person`.`name`;
    ```
3. 为了区别不同 Tag（label） 的属性，使用属性时必须额外指定Tag名称。即<变量名>.<Tag名>.<属性名>，例如：
    - 访问business_system节点的name属性：`bs.business_system.name`
    - 访问database节点的name属性：`db.database.name`
    - 错误示例：`db.name`（缺少Tag名）
    - 正确示例：`db.database.name`（完整Tag+属性名）
4. EDGE（relationship）的使用和标准Cypher一致，使用EDGE属性时不需指定EDGE名称。即<变量名>.<属性名>，例如s.name。
5. 如果生成的cypher语句中有查询条件，查询条件中请遵循==是相等运算符，=是赋值运算符的规则。
6. 当要对指定的值进行计数并返回结果时，请直接使用COUNT()函数，并通过RETURN语句输出计数值，图数据库不需要使用GROUP BY来进行聚合操作，请参考以下示例：
    - 查询文件的数据级别分布情况：
    ```
    MATCH (f:`file`)
    RETURN f.`file`.data_level as data_level, COUNT(f) as amount 
    order by amount desc
    ```
7. 当用户发起查询请求时，请首先判断用户是否指定了具体的查询主体（例如特定的表table、字段field或文件file）。
    - 如果用户未指定具体主体，则查询应专注于带有data_directory标签的数据，以获取分类分级汇总结果并进行统计分析。此时，不应包含其他类型的标签（如business_system、database、table、field或file）在查询中。
    - 如果用户确实指定了具体主体，则应针对该主体对应的标签（如business_system、database、table、field或file），执行查询以获得其特定的分类分级结果。
   请根据用户的输入动态调整查询逻辑，确保只查询相关的标签，以提供最相关的结果给用户。
8. 如果用户的问题需要基于两个不同的标签（tag）来检索数据，为了正确地合并这两个结果集，请确保使用 UNION 操作符，并且在进行 UNION 操作时，请使用as关键字，给字句返回的列设置相同的别名，请参考以下示例：
    - 查询业务系统、数据库和文件系统的数量：
    ```
    MATCH (bs:`business_system`)
    RETURN count(bs) as amount
    UNION
    MATCH (db:`database`)
    RETURN count(db) as amount
    UNION
    MATCH (fs:`file_system`)
    RETURN count(fs) as amount
    ```
9. 如果用户的问题之前已经问过并得到回答，请准确重复之前的答案。
10. 仅输出cypher语句，不要输出任何其他不相关内容，不要使用markdown格式输出。
"""
DEFAULT_NEBULAGRAPH_NL2CYPHER_PROMPT = PromptTemplate(
    DEFAULT_NEBULAGRAPH_NL2CYPHER_PROMPT_TMPL1,
    prompt_type=PromptType.TEXT_TO_GRAPH_QUERY,
)

# Prompt
DEFAULT_NEO4J_NL2CYPHER_PROMPT_TMPL = (
    "Task:Generate Cypher statement to query a graph database.\n"
    "Instructions:\n"
    "Use only the provided relationship types and properties in the schema.\n"
    "Do not use any other relationship types or properties that are not provided.\n"
    "Schema:\n"
    "{schema}\n"
    "Note: Do not include any explanations or apologies in your responses.\n"
    "Do not respond to any questions that might ask anything else than for you "
    "to construct a Cypher statement. \n"
    "Do not include any text except the generated Cypher statement.\n"
    "\n"
    "The question is:\n"
    "{query_str}\n"
)

DEFAULT_NEO4J_NL2CYPHER_PROMPT = PromptTemplate(
    DEFAULT_NEO4J_NL2CYPHER_PROMPT_TMPL,
    prompt_type=PromptType.TEXT_TO_GRAPH_QUERY,
)

DEFAULT_KG_RESPONSE_ANSWER_PROMPT_TMPL = """
The original question is given below.
This question has been translated into a Graph Database query.
Both the Graph query and the response are given below.
Given the Graph Query response, synthesise a response to the original question.

Original question: {query_str}
Graph query: {kg_query_str}
Graph response: {kg_response_str}
Response:
"""

DEFAULT_KG_RESPONSE_ANSWER_PROMPT_TMPL1 = """
原始问题如下。
此问题已转换为图形数据库查询。
下面给出了图数据库的查询和响应。
你作为一个数据分类分级专家，根据原始问题、图数据库查询和响应，综合分析生成对原始问题的答案，并对这些数据进行数据分类分级与数据安全方面的进一步分析。

原始问题: {query_str}
图形查询: {kg_query_str}
图形响应: {kg_response_str}
答复：
"""

DEFAULT_KG_RESPONSE_ANSWER_PROMPT_TMPL2 = """
## Role  
你是一个智能数据分析助手，请根据上下文，回答用户的问题，最后生成3个自然连贯的后续问题。

## Goals  
使用Schema中NebulaGraph图数据库结构，Context中的用户原始问题、图谱查询语句、查询结果集以及相关知识，回答用户的问题，最后生成3个后续问题列表，回答用户的问题和生成后续问题必须严格遵守Attention中的规则。

## Schema  
[图数据库schema]
{schema}  

## Context  
[用户原始问题]  
{query_str} 

[图谱查询语句]  
{kg_query_str} 

[查询结果集]  
{kg_response_str}

[相关知识]

## Attention    
1. 严格限定分析范围在schema定义字段内  
2. 使用[关键结论]->[数据分析]->[治理建议]的三段式结构，禁用模糊表述，所有结论需标注数据支撑点
3. [关键结论]使用结构化呈现（对比表格/列表），能使用表格则使用表格展示
4. [数据分析]除了分析数据本身，请结合数据分类分级安全做进一步分析
5. [治理建议]根据用户问题灵活从以下几个方面范围内来给出，且要与数据进行结合，给出具体建议，禁用模糊表述：
   - 数据保护策略
   - 风险管理与合规性
   - 数据生命周期管理
   - 技术与工具支持
   - 培训与文化建设
   - 监控与优化
6. 后续问题的生成需遵循以下原则：
   - 相关性：问题必须基于当前查询结果或Schema中关联的实体/关系扩展
   - 多样性：覆盖不同分析角度（如因果推理、对比、趋势预测、异常挖掘、假设场景）
   - 具体性：避免笼统提问，需包含Schema中的具体属性（如时间、地点、指标名称）
   - 唯一性：避免与用户原始问题重复 
   - 稳定性：后续问题应保证能够使用NebulaGraph图数据库查询语句在schema的范围内查询出结果，如果不能确定问题对应的查询语句是否一定能够正常查出结果，则请简化问题。
7. 生成的3个后续问题要包含在<follow-questions></follow-questions>标签内，每行一个。例如：
<follow-questions>>
1. 统计数据级别的分布情况。
2. 统计统计数据类别Top 10。
3. 统计有多少业务系统、数据库和文件系统？
</follow-questions>>
"""

DEFAULT_KG_RESPONSE_ANSWER_PROMPT = PromptTemplate(
    DEFAULT_KG_RESPONSE_ANSWER_PROMPT_TMPL2,
    prompt_type=PromptType.QUESTION_ANSWER,
)

DEFAULT_TEXT_QA_PROMPT_TMPL = (
    "{context_str}"
)
DEFAULT_TEXT_QA_PROMPT = PromptTemplate(
    DEFAULT_TEXT_QA_PROMPT_TMPL, prompt_type=PromptType.QUESTION_ANSWER
)


def format_query(content: str):
    # 使用正则表达式匹配```json ... ```模式
    pattern = r"```cypher(.*?)```"
    matches = re.findall(pattern, content, re.DOTALL)

    if matches:
        # 假设只有一个匹配项
        content = matches[0].strip()
    return content


class MyKGQueryEngine(KnowledgeGraphQueryEngine):
    def __init__(self, llm, stream=True, *args, **kwargs):
        self._max_tokens = kwargs.get("max_tokens", 20000)
        callback_manager = Settings.callback_manager
        response_synthesizer = get_response_synthesizer(
            llm=llm,
            text_qa_template=DEFAULT_TEXT_QA_PROMPT,
            callback_manager=callback_manager,
            streaming=stream,
            verbose=True
        )
        super().__init__(
            llm=llm,
            graph_query_synthesis_prompt=DEFAULT_NEBULAGRAPH_NL2CYPHER_PROMPT,
            graph_response_answer_prompt=DEFAULT_KG_RESPONSE_ANSWER_PROMPT,
            response_synthesizer=response_synthesizer,
            *args,
            **kwargs
        )

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Get nodes for response."""
        graph_store_query = self.generate_query(query_bundle.query_str)
        if self._verbose:
            print_text(f"Graph Store Query:\n{graph_store_query}\n", color="yellow")
        logger.debug(f"Graph Store Query:\n{graph_store_query}")

        with self.callback_manager.event(
                CBEventType.RETRIEVE,
                payload={EventPayload.QUERY_STR: graph_store_query},
        ) as retrieve_event:
            # Get the graph store response
            graph_store_response = self.graph_store.query(query=graph_store_query)
            if self._verbose:
                print_text(
                    f"Graph Store Response:\n{graph_store_response}\n",
                    color="yellow",
                )
            logger.debug(f"Graph Store Response:\n{graph_store_response}")

            retrieve_event.on_end(payload={EventPayload.RESPONSE: graph_store_response})

        retrieved_graph_context: Sequence = self._graph_response_answer_prompt.format(
            query_str=query_bundle.query_str,
            kg_query_str=graph_store_query,
            kg_response_str=graph_store_response,
            schema=self._graph_schema,
        )

        node = NodeWithScore(
            node=TextNode(
                text=retrieved_graph_context,
                metadata={
                    # "query_str": query_bundle.query_str,
                    # "graph_store_query": graph_store_query,
                    # "graph_store_response": graph_store_response,
                    # "graph_schema": self._graph_schema,
                },
            ),
            score=1.0,
        )
        return [node]

    def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
        """Query the graph store."""
        with self.callback_manager.event(
            CBEventType.QUERY, payload={EventPayload.QUERY_STR: query_bundle.query_str}
        ) as query_event:
            nodes: List[NodeWithScore] = self._retrieve(query_bundle)

            response = self._response_synthesizer.synthesize(
                query=query_bundle,
                nodes=nodes,
            )

            # for chunk in response.response_gen:
            #     print_text(chunk, color="green")

            # if self._verbose:
            #     print_text(f"Final Response: {response}\n", color="green")

            query_event.on_end(payload={EventPayload.RESPONSE: response})

        return response

    def generate_query(self, query_str: str) -> str:
        """Generate a Graph Store Query from a query bundle."""
        # Get the query engine query string

        question_sql_list = vn.get_similar_question_sql(query_str)
        ddl_list = vn.get_related_ddl(query_str)
        doc_list = vn.get_related_documentation(query_str)
        # 计算当前时间
        now = datetime.now()
        current_time = now.strftime("%Y-%m-%d %H:%M:%S")
        initial_prompt = self._graph_query_synthesis_prompt.format(current_time=current_time)

        messages = self.get_sql_prompt(
            initial_prompt=initial_prompt,
            question=query_str,
            question_sql_list=question_sql_list,
            ddl_list=ddl_list,
            doc_list=doc_list,
            schema=self._graph_schema
        )

        # graph_store_query: str = self._llm.predict(
        #     self._graph_query_synthesis_prompt,
        #     query_str=query_str,
        #     schema=self._graph_schema
        # )
        chat_response = self._llm.chat(
            messages=messages
        )

        graph_store_query = chat_response.message.content or ""

        graph_store_query = format_query(graph_store_query)

        return graph_store_query

    def get_sql_prompt(
            self,
            initial_prompt: str,
            question: str,
            question_sql_list: list,
            ddl_list: list,
            doc_list: list,
            schema: str
    ):
        initial_prompt = self.add_ddl_to_prompt(
            initial_prompt=initial_prompt,
            ddl_list=ddl_list,
            schema=schema,
            max_tokens=self._max_tokens,
        )
        initial_prompt = self.add_documentation_to_prompt(
            initial_prompt=initial_prompt,
            documentation_list=doc_list,
            max_tokens=self._max_tokens,
        )
        messages = [ChatMessage(role=MessageRole.SYSTEM, content=initial_prompt)]
        for question_sql in question_sql_list:
            if question_sql is None:
                print("question_sql is None")
            else:
                if question_sql is not None and "question" in question_sql and "sql" in question_sql:
                    messages.append(ChatMessage(role=MessageRole.USER, content=question_sql["question"]))
                    messages.append(ChatMessage(role=MessageRole.ASSISTANT, content=question_sql["sql"]))

        messages.append(ChatMessage(role=MessageRole.USER, content=question))
        return messages

    def add_ddl_to_prompt(
            self, initial_prompt: str, ddl_list: list[str], schema: str, max_tokens: int = 20000
    ) -> str:
        if len(ddl_list) > 0:
            for ddl in ddl_list:
                if (
                        self.str_to_approx_token_count(initial_prompt)
                        + self.str_to_approx_token_count(schema)
                        + self.str_to_approx_token_count(ddl)
                        < max_tokens
                ):
                    schema += f"{ddl}\n\n"

        initial_prompt = PromptTemplate(initial_prompt, prompt_type=PromptType.TEXT_TO_GRAPH_QUERY).format(schema=schema)
        return initial_prompt

    def add_documentation_to_prompt(
            self,
            initial_prompt: str,
            documentation_list: list[str],
            max_tokens: int = 20000,
    ) -> str:
        if len(documentation_list) > 0:
            context = ""

            for documentation in documentation_list:
                if (
                        self.str_to_approx_token_count(initial_prompt)
                        + self.str_to_approx_token_count(context)
                        + self.str_to_approx_token_count(documentation)
                        < max_tokens
                ):
                    context += f"{documentation}\n\n"

        initial_prompt = PromptTemplate(initial_prompt, prompt_type=PromptType.TEXT_TO_GRAPH_QUERY).format(context=context)
        return initial_prompt

    def str_to_approx_token_count(self, string: str) -> int:
        return len(string) / 4
